#include <crt_util.h>

void* avl_walk(avl_tree_t* tree, void* oldnode, int left)
{
    size_t off = tree->avl_offset;
    avl_node_t* node = AVL_DATA2NODE(oldnode, off);
    int right = 1 - left;
    int was_child;

    if (node == NULL)
        return (NULL);

    if (node->avl_child[left] != NULL) {
        for (node = node->avl_child[left]; node->avl_child[right] != NULL; node = node->avl_child[right])
            ;

    } else {
        for (;;) {
            was_child = AVL_XCHILD(node);
            node = AVL_XPARENT(node);
            if (node == NULL)
                return (NULL);
            if (was_child == right)
                break;
        }
    }

    return (AVL_NODE2DATA(node, off));
}

void* avl_first(avl_tree_t* tree)
{
    avl_node_t* node;
    avl_node_t* prev = NULL;
    size_t off = tree->avl_offset;

    for (node = tree->avl_root; node != NULL; node = node->avl_child[0])
        prev = node;

    if (prev != NULL)
        return (AVL_NODE2DATA(prev, off));
    return (NULL);
}

void* avl_last(avl_tree_t* tree)
{
    avl_node_t* node;
    avl_node_t* prev = NULL;
    size_t off = tree->avl_offset;

    for (node = tree->avl_root; node != NULL; node = node->avl_child[1])
        prev = node;

    if (prev != NULL)
        return (AVL_NODE2DATA(prev, off));
    return (NULL);
}

void* avl_nearest(avl_tree_t* tree, avl_index_t where, int direction)
{
    int child = AVL_INDEX2CHILD(where);
    avl_node_t* node = AVL_INDEX2NODE(where);
    void* data;
    size_t off = tree->avl_offset;

    if (node == NULL) {
        return (NULL);
    }
    data = AVL_NODE2DATA(node, off);
    if (child != direction)
        return (data);

    return (avl_walk(tree, data, direction));
}

/*
 * Search for the node which contains "value".  The algorithm is a
 * simple binary tree search.
 *
 * return value:
 *	NULL: the value is not in the AVL tree
 *		*where (if not NULL)  is set to indicate the insertion point
 *	"void *"  of the found tree node
 */
void* avl_find(avl_tree_t* tree, void* value, avl_index_t* where)
{
    avl_node_t* node;
    avl_node_t* prev = NULL;
    int child = 0;
    int diff;
    size_t off = tree->avl_offset;
    int avl_balance2child[] = {0, 0, 1};

    for (node = tree->avl_root; node != NULL; node = node->avl_child[child]) {
        prev = node;

        diff = tree->avl_compar(value, AVL_NODE2DATA(node, off));
        if (diff == 0) {
#ifdef DEBUG
            if (where != NULL)
                *where = 0;
#endif
            return (AVL_NODE2DATA(node, off));
        }
        child = avl_balance2child[1 + diff];
    }

    if (where != NULL)
        *where = AVL_MKINDEX(prev, child);

    return (NULL);
}

/*
 * Perform a rotation to restore balance at the subtree given by depth.
 *
 * This routine is used by both insertion and deletion. The return value
 * indicates:
 *	 0 : subtree did not change height
 *	!0 : subtree was reduced in height
 *
 * The code is written as if handling left rotations, right rotations are
 * symmetric and handled by swapping values of variables right/left[_heavy]
 *
 * On input balance is the "new" balance at "node". This value is either
 * -2 or +2.
 */
static int avl_rotation(avl_tree_t* tree, avl_node_t* node, int balance)
{
    int left = !(balance < 0); /* when balance = -2, left will be 0 */
    int right = 1 - left;
    int left_heavy = balance >> 1;
    int right_heavy = -left_heavy;
    avl_node_t* parent = AVL_XPARENT(node);
    avl_node_t* child = node->avl_child[left];
    avl_node_t* cright;
    avl_node_t* gchild;
    avl_node_t* gright;
    avl_node_t* gleft;
    int which_child = AVL_XCHILD(node);
    int child_bal = AVL_XBALANCE(child);

    /* BEGIN CSTYLED */
    /*
     * case 1 : node is overly left heavy, the left child is balanced or
     * also left heavy. This requires the following rotation.
     *
     *                   (node bal:-2)
     *                    /           \
     *                   /             \
     *              (child bal:0 or -1)
     *              /    \
     *             /      \
     *                     cright
     *
     * becomes:
     *
     *              (child bal:1 or 0)
     *              /        \
     *             /          \
     *                        (node bal:-1 or 0)
     *                         /     \
     *                        /       \
     *                     cright
     *
     * we detect this situation by noting that child's balance is not
     * right_heavy.
     */
    /* END CSTYLED */
    if (child_bal != right_heavy) {
        /*
         * compute new balance of nodes
         *
         * If child used to be left heavy (now balanced) we reduced
         * the height of this sub-tree -- used in "return...;" below
         */
        child_bal += right_heavy; /* adjust towards right */

        /*
         * move "cright" to be node's left child
         */
        cright = child->avl_child[right];
        node->avl_child[left] = cright;
        if (cright != NULL) {
            AVL_SETPARENT(cright, node);
            AVL_SETCHILD(cright, left);
        }

        /*
         * move node to be child's right child
         */
        child->avl_child[right] = node;
        AVL_SETBALANCE(node, -child_bal);
        AVL_SETCHILD(node, right);
        AVL_SETPARENT(node, child);

        /*
         * update the pointer into this subtree
         */
        AVL_SETBALANCE(child, child_bal);
        AVL_SETCHILD(child, which_child);
        AVL_SETPARENT(child, parent);
        if (parent != NULL)
            parent->avl_child[which_child] = child;
        else
            tree->avl_root = child;

        return (child_bal == 0);
    }

    /* BEGIN CSTYLED */
    /*
     * case 2 : When node is left heavy, but child is right heavy we use
     * a different rotation.
     *
     *                   (node b:-2)
     *                    /   \
     *                   /     \
     *                  /       \
     *             (child b:+1)
     *              /     \
     *             /       \
     *                   (gchild b: != 0)
     *                     /  \
     *                    /    \
     *                 gleft   gright
     *
     * becomes:
     *
     *              (gchild b:0)
     *              /       \
     *             /         \
     *            /           \
     *        (child b:?)   (node b:?)
     *         /  \          /   \
     *        /    \        /     \
     *            gleft   gright
     *
     * computing the new balances is more complicated. As an example:
     *	 if gchild was right_heavy, then child is now left heavy
     *		else it is balanced
     */
    /* END CSTYLED */
    gchild = child->avl_child[right];
    gleft = gchild->avl_child[left];
    gright = gchild->avl_child[right];

    /*
     * move gright to left child of node and
     *
     * move gleft to right child of node
     */
    node->avl_child[left] = gright;
    if (gright != NULL) {
        AVL_SETPARENT(gright, node);
        AVL_SETCHILD(gright, left);
    }

    child->avl_child[right] = gleft;
    if (gleft != NULL) {
        AVL_SETPARENT(gleft, child);
        AVL_SETCHILD(gleft, right);
    }

    /*
     * move child to left child of gchild and
     *
     * move node to right child of gchild and
     *
     * fixup parent of all this to point to gchild
     */
    balance = AVL_XBALANCE(gchild);
    gchild->avl_child[left] = child;
    AVL_SETBALANCE(child, (balance == right_heavy ? left_heavy : 0));
    AVL_SETPARENT(child, gchild);
    AVL_SETCHILD(child, left);

    gchild->avl_child[right] = node;
    AVL_SETBALANCE(node, (balance == left_heavy ? right_heavy : 0));
    AVL_SETPARENT(node, gchild);
    AVL_SETCHILD(node, right);

    AVL_SETBALANCE(gchild, 0);
    AVL_SETPARENT(gchild, parent);
    AVL_SETCHILD(gchild, which_child);
    if (parent != NULL)
        parent->avl_child[which_child] = gchild;
    else
        tree->avl_root = gchild;

    return (1); /* the new tree is always shorter */
}

/*
 * Insert a new node into an AVL tree at the specified (from avl_find()) place.
 *
 * Newly inserted nodes are always leaf nodes in the tree, since avl_find()
 * searches out to the leaf positions.  The avl_index_t indicates the node
 * which will be the parent of the new node.
 *
 * After the node is inserted, a single rotation further up the tree may
 * be necessary to maintain an acceptable AVL balance.
 */
void avl_insert(avl_tree_t* tree, void* new_data, avl_index_t where)
{
    rc_return(tree != NULL);
    avl_node_t* node;
    avl_node_t* parent = AVL_INDEX2NODE(where);
    int old_balance;
    int new_balance;
    int which_child = AVL_INDEX2CHILD(where);
    size_t off = tree->avl_offset;

    int avl_child2balance[2] = {-1, 1};

    node = AVL_DATA2NODE(new_data, off);

    /*
     * First, add the node to the tree at the indicated position.
     */
    ++tree->avl_numnodes;

    node->avl_child[0] = NULL;
    node->avl_child[1] = NULL;

    AVL_SETCHILD(node, which_child);
    AVL_SETBALANCE(node, 0);
    AVL_SETPARENT(node, parent);
    if (parent != NULL) {
        parent->avl_child[which_child] = node;
    } else {
        tree->avl_root = node;
    }
    /*
     * Now, back up the tree modifying the balance of all nodes above the
     * insertion point. If we get to a highly unbalanced ancestor, we
     * need to do a rotation.  If we back out of the tree we are done.
     * If we brought any subtree into perfect balance (0), we are also done.
     */
    for (;;) {
        node = parent;
        if (node == NULL)
            return;

        /*
         * Compute the new balance
         */
        old_balance = AVL_XBALANCE(node);
        new_balance = old_balance + avl_child2balance[which_child];

        /*
         * If we introduced equal balance, then we are done immediately
         */
        if (new_balance == 0) {
            AVL_SETBALANCE(node, 0);
            return;
        }

        /*
         * If both old and new are not zero we went
         * from -1 to -2 balance, do a rotation.
         */
        if (old_balance != 0)
            break;

        AVL_SETBALANCE(node, new_balance);
        parent = AVL_XPARENT(node);
        which_child = AVL_XCHILD(node);
    }

    /*
     * perform a rotation to fix the tree and return
     */
    (void)avl_rotation(tree, node, new_balance);
}

/*
 * Insert "new_data" in "tree" in the given "direction" either after or
 * before (AVL_AFTER, AVL_BEFORE) the data "here".
 *
 * Insertions can only be done at empty leaf points in the tree, therefore
 * if the given child of the node is already present we move to either
 * the AVL_PREV or AVL_NEXT and reverse the insertion direction. Since
 * every other node in the tree is a leaf, this always works.
 *
 * To help developers using this interface, we assert that the new node
 * is correctly ordered at every step of the way in DEBUG kernels.
 */
void avl_insert_here(avl_tree_t* tree, void* new_data, void* here, int direction)
{
    avl_node_t* node;
    int child = direction; /* rely on AVL_BEFORE == 0, AVL_AFTER == 1 */

    rc_return(tree != NULL);
    rc_return(new_data != NULL);
    rc_return(here != NULL);
    rc_return(direction == AVL_BEFORE || direction == AVL_AFTER);

    /*
     * If corresponding child of node is not NULL, go to the neighboring
     * node and reverse the insertion direction.
     */
    node = AVL_DATA2NODE(here, tree->avl_offset);

    if (node->avl_child[child] != NULL) {
        node = node->avl_child[child];
        child = 1 - child;
        while (node->avl_child[child] != NULL) {
            node = node->avl_child[child];
        }
    }
    avl_insert(tree, new_data, AVL_MKINDEX(node, child));
}

/*
 * Add a new node to an AVL tree.
 */
void avl_add(avl_tree_t* tree, void* new_node)
{
    avl_index_t where = 0;
    avl_insert(tree, new_node, where);
}

/*
 * Delete a node from the AVL tree.  Deletion is similar to insertion, but
 * with 2 complications.
 *
 * First, we may be deleting an interior node. Consider the following subtree:
 *
 *     d           c            c
 *    / \         / \          / \
 *   b   e       b   e        b   e
 *  / \	        / \          /
 * a   c       a            a
 *
 * When we are deleting node (d), we find and bring up an adjacent valued leaf
 * node, say (c), to take the interior node's place. In the code this is
 * handled by temporarily swapping (d) and (c) in the tree and then using
 * common code to delete (d) from the leaf position.
 *
 * Secondly, an interior deletion from a deep tree may require more than one
 * rotation to fix the balance. This is handled by moving up the tree through
 * parents and applying rotations as needed. The return value from
 * avl_rotation() is used to detect when a subtree did not change overall
 * height due to a rotation.
 */
void avl_remove(avl_tree_t* tree, void* data)
{
    avl_node_t* delete;
    avl_node_t* parent;
    avl_node_t* node;
    avl_node_t tmp;
    int old_balance;
    int new_balance;
    int left;
    int right;
    int which_child;
    size_t off = tree->avl_offset;

    int avl_child2balance[2] = {-1, 1};
    int avl_balance2child[] = {0, 0, 1};

    rc_return(tree != NULL);

    delete = AVL_DATA2NODE(data, off);

    /*
     * Deletion is easiest with a node that has at most 1 child.
     * We swap a node with 2 children with a sequentially valued
     * neighbor node. That node will have at most 1 child. Note this
     * has no effect on the ordering of the remaining nodes.
     *
     * As an optimization, we choose the greater neighbor if the tree
     * is right heavy, otherwise the left neighbor. This reduces the
     * number of rotations needed.
     */
    if (delete->avl_child[0] != NULL && delete->avl_child[1] != NULL) {
        /*
         * choose node to swap from whichever side is taller
         */
        old_balance = AVL_XBALANCE(delete);
        left = avl_balance2child[old_balance + 1];
        right = 1 - left;

        /*
         * get to the previous value'd node
         * (down 1 left, as far as possible right)
         */
        for (node = delete->avl_child[left]; node->avl_child[right] != NULL; node = node->avl_child[right])
            ;

        /*
         * create a temp placeholder for 'node'
         * move 'node' to delete's spot in the tree
         */
        tmp = *node;

        *node = *delete;
        if (node->avl_child[left] == node)
            node->avl_child[left] = &tmp;

        parent = AVL_XPARENT(node);
        if (parent != NULL)
            parent->avl_child[AVL_XCHILD(node)] = node;
        else
            tree->avl_root = node;
        AVL_SETPARENT(node->avl_child[left], node);
        AVL_SETPARENT(node->avl_child[right], node);

        /*
         * Put tmp where node used to be (just temporary).
         * It always has a parent and at most 1 child.
         */
        delete = &tmp;
        parent = AVL_XPARENT(delete);
        parent->avl_child[AVL_XCHILD(delete)] = delete;
        which_child = (delete->avl_child[1] != 0);
        if (delete->avl_child[which_child] != NULL)
            AVL_SETPARENT(delete->avl_child[which_child], delete);
    }

    /*
     * Here we know "delete" is at least partially a leaf node. It can
     * be easily removed from the tree.
     */
    rc_return(tree->avl_numnodes > 0);

    --tree->avl_numnodes;
    parent = AVL_XPARENT(delete);
    which_child = AVL_XCHILD(delete);
    if (delete->avl_child[0] != NULL)
        node = delete->avl_child[0];
    else
        node = delete->avl_child[1];

    /*
     * Connect parent directly to node (leaving out delete).
     */
    if (node != NULL) {
        AVL_SETPARENT(node, parent);
        AVL_SETCHILD(node, which_child);
    }
    if (parent == NULL) {
        tree->avl_root = node;
        return;
    }
    parent->avl_child[which_child] = node;

    /*
     * Since the subtree is now shorter, begin adjusting parent balances
     * and performing any needed rotations.
     */
    do {
        /*
         * Move up the tree and adjust the balance
         *
         * Capture the parent and which_child values for the next
         * iteration before any rotations occur.
         */
        node = parent;
        old_balance = AVL_XBALANCE(node);
        new_balance = old_balance - avl_child2balance[which_child];
        parent = AVL_XPARENT(node);
        which_child = AVL_XCHILD(node);

        /*
         * If a node was in perfect balance but isn't anymore then
         * we can stop, since the height didn't change above this point
         * due to a deletion.
         */
        if (old_balance == 0) {
            AVL_SETBALANCE(node, new_balance);
            break;
        }

        /*
         * If the new balance is zero, we don't need to rotate
         * else
         * need a rotation to fix the balance.
         * If the rotation doesn't change the height
         * of the sub-tree we have finished adjusting.
         */
        if (new_balance == 0)
            AVL_SETBALANCE(node, new_balance);
        else if (!avl_rotation(tree, node, new_balance))
            break;
    } while (parent != NULL);
}

#define AVL_REINSERT(tree, obj) \
    avl_remove((tree), (obj)); \
    avl_add((tree), (obj))

int avl_update_lt(avl_tree_t* t, void* obj)
{
    void* neighbor;

    neighbor = AVL_PREV(t, obj);
    if ((neighbor != NULL) && (t->avl_compar(obj, neighbor) < 0)) {
        AVL_REINSERT(t, obj);
        return S_SUCCESS;
    }

    return S_ERROR;
}

int avl_update_gt(avl_tree_t* t, void* obj)
{
    void* neighbor;

    neighbor = AVL_NEXT(t, obj);
    if ((neighbor != NULL) && (t->avl_compar(obj, neighbor) > 0)) {
        AVL_REINSERT(t, obj);
        return S_SUCCESS;
    }

    return S_ERROR;
}

int avl_update(avl_tree_t* t, void* obj)
{
    void* neighbor;

    neighbor = AVL_PREV(t, obj);
    if ((neighbor != NULL) && (t->avl_compar(obj, neighbor) < 0)) {
        AVL_REINSERT(t, obj);
        return S_SUCCESS;
    }

    neighbor = AVL_NEXT(t, obj);
    if ((neighbor != NULL) && (t->avl_compar(obj, neighbor) > 0)) {
        AVL_REINSERT(t, obj);
        return S_SUCCESS;
    }

    return S_ERROR;
}

/*
 * initialize a new AVL tree
 */
void avl_create(avl_tree_t* tree, int (*compar)(const void*, const void*), size_t size, size_t offset)
{
    rc_return(tree != NULL);
    rc_return(compar != NULL);
    rc_return(size > 0);
    rc_return(size >= offset + sizeof(avl_node_t));

    tree->avl_compar = compar;
    tree->avl_root = NULL;
    tree->avl_numnodes = 0;
    tree->avl_size = size;
    tree->avl_offset = offset;
}

void avl_destroy(avl_tree_t* tree)
{
}

unsigned long avl_numnodes(avl_tree_t* tree)
{
    rc_error(tree != NULL, 0);
    return (tree->avl_numnodes);
}

int avl_is_empty(avl_tree_t* tree)
{
    rc_error(tree != NULL, S_ERROR);
    return (tree->avl_numnodes == 0 ? S_SUCCESS : S_ERROR);
}

#define CHILDBIT (1L)

void* avl_destroy_nodes(avl_tree_t* tree, void** cookie)
{
    avl_node_t* node;
    avl_node_t* parent;
    int child;
    void* first;
    size_t off = tree->avl_offset;

    /*
     * Initial calls go to the first node or it's right descendant.
     */
    if (*cookie == NULL) {
        first = avl_first(tree);

        /*
         * deal with an empty tree
         */
        if (first == NULL) {
            *cookie = (void*)CHILDBIT;
            return (NULL);
        }

        node = AVL_DATA2NODE(first, off);
        parent = AVL_XPARENT(node);
        goto check_right_side;
    }

    /*
     * If there is no parent to return to we are done.
     */
    parent = (avl_node_t*)((uintptr_t)(*cookie) & ~CHILDBIT);
    if (parent == NULL) {
        if (tree->avl_root != NULL) {
            tree->avl_root = NULL;
            tree->avl_numnodes = 0;
        }
        return (NULL);
    }

    /*
     * Remove the child pointer we just visited from the parent and tree.
     */
    child = (uintptr_t)(*cookie) & CHILDBIT;
    parent->avl_child[child] = NULL;

    --tree->avl_numnodes;

    /*
     * If we just did a right child or there isn't one, go up to parent.
     */
    if (child == 1 || parent->avl_child[1] == NULL) {
        node = parent;
        parent = AVL_XPARENT(parent);
        goto done;
    }

    /*
     * Do parent's right child, then leftmost descendent.
     */
    node = parent->avl_child[1];
    while (node->avl_child[0] != NULL) {
        parent = node;
        node = node->avl_child[0];
    }

    /*
     * If here, we moved to a left child. It may have one
     * child on the right (when balance == +1).
     */
check_right_side:
    if (node->avl_child[1] != NULL) {
        parent = node;
        node = node->avl_child[1];

    } else {
    }

done:
    if (parent == NULL) {
        *cookie = (void*)CHILDBIT;
    } else {
        *cookie = (void*)((uintptr_t)parent | AVL_XCHILD(node));
    }

    return (AVL_NODE2DATA(node, off));
}
